# JPEG 图像压缩
# 变量类型声明
from __future__ import annotations

import cv2 as cv
import numpy as np

# 亮度量化表
L1 = np.array([
    [17, 18, 24, 99, 99, 99, 99, 99],
    [18, 21, 26, 66, 99, 99, 99, 99],
    [24, 26, 56, 99, 99, 99, 99, 99],
    [47, 66, 99, 99, 99, 99, 99, 99],
    [99, 99, 99, 99, 99, 99, 99, 99],
    [99, 99, 99, 99, 99, 99, 99, 99],
    [99, 99, 99, 99, 99, 99, 99, 99],
    [99, 99, 99, 99, 99, 99, 99, 99],
], dtype=np.float64)

# 色度量化表
L2 = np.array([
    [16, 11, 10, 16, 24, 40, 51, 61],
    [12, 12, 14, 19, 26, 58, 60, 55],
    [14, 13, 16, 24, 40, 57, 69, 56],
    [14, 17, 22, 29, 51, 87, 80, 62],
    [18, 22, 37, 56, 68, 109, 103, 77],
    [24, 35, 55, 64, 81, 104, 113, 92],
    [49, 64, 78, 87, 103, 121, 120, 101],
    [72, 92, 95, 98, 112, 100, 103, 99],
], dtype=np.float64)

# 游程编码表
Z = np.array([
    [0, 1, 5, 6, 14, 15, 27, 28],
    [2, 4, 7, 13, 16, 26, 29, 42],
    [3, 8, 12, 17, 25, 30, 41, 43],
    [9, 11, 18, 24, 31, 40, 44, 53],
    [10, 19, 23, 32, 39, 45, 52, 54],
    [20, 22, 33, 38, 46, 51, 55, 60],
    [21, 34, 37, 47, 50, 56, 59, 61],
    [35, 36, 48, 49, 57, 58, 62, 63],
], dtype=np.int32)


# 游程编码
def DLE(zzag_x):
    ls = []
    tmp = None
    cnt = 0
    for i in zzag_x:
        if tmp is not None:
            if tmp == i:
                cnt += 1
            else:
                ls.append((cnt, tmp))
                tmp = i
                cnt = 1
        else:
            tmp = i
            cnt = 1
    ls.append((cnt, tmp))
    return ls


def rDLE(dle: [(int, int)]) -> np.array:
    size = 0
    for i in dle:
        size += i[0]
    zzag_x = np.ones(size, dtype=np.int32)
    cnt = 0
    for i in dle:
        for j in range(i[0]):
            zzag_x[cnt] = i[1]
            cnt += 1
    return zzag_x


# 生成 (游程编码符号, 频数) 表
def gen_freq_table(ls: list) -> [((int, int), int)]:
    els = sorted(list(set(ls)), key=ls.index)
    return list(map(lambda x: (x, ls.count(x)), els))


# Huffman 树节点
class Node:

    def __init__(self):
        # 左子树
        self.left: Node | None = None
        # 右子树
        self.right: Node | None = None
        # 游程编码的符号
        self.sign: (int, int) = None
        # 该树的总频数
        self.freq: int = 0


# 生成哈夫曼树
def create_tree_from_freq_table(freq_table: dict) -> [Node]:
    tree_holder = list()

    # 先把所有 符号-频数 对 转为树节点
    for i in freq_table:
        node = Node()
        node.sign = i[0]
        node.freq = i[1]
        tree_holder.append(node)

    # 只有一个节点则放在左子树上
    if len(tree_holder) == 1:
        node = Node()
        node.left = tree_holder[0]
        tree_holder[0] = node

    # 构造 huffman 树
    while len(tree_holder) != 1:
        leastNode1: Node = Node()
        leastNode2: Node = Node()
        leastNode1.freq = (-1 >> 1) & 0x7FFFFFFF
        leastNode2.freq = leastNode1.freq
        for i in tree_holder:
            if i.freq < leastNode1.freq:
                leastNode1 = i
        tree_holder.remove(leastNode1)
        for i in tree_holder:
            if i.freq < leastNode2.freq:
                leastNode2 = i
        tree_holder.remove(leastNode2)
        node = Node()
        node.left = leastNode1
        node.right = leastNode2
        node.freq = leastNode1.freq + leastNode2.freq
        tree_holder.insert(0, node)
    return tree_holder


# 深度优先遍历生成编码表
def dfs(root: Node | None, table: dict, current_code: str):
    if root is None:
        return
    if root.sign is not None:
        table[root.sign] = current_code
    dfs(root.left, table, current_code + '0')
    dfs(root.right, table, current_code + '1')


def huffman_encode(dle_x: np.array) -> (str, Node):
    freq_table = gen_freq_table(dle_x)
    ls = create_tree_from_freq_table(freq_table)
    d = dict()
    dfs(ls[0], d, '')
    code_x = ''
    for i in dle_x:
        code_x += d[i]
    return code_x, ls[0]


# 哈夫曼解码函数
def huffman_decode(huffman_code: str, tree: Node) -> [(int, int)]:
    decode_dle_x = list()
    tmp = tree
    for i in huffman_code:
        if i == '0':
            tmp = tmp.left
        elif i == '1':
            tmp = tmp.right
        if tmp.sign is not None:
            decode_dle_x.append(tmp.sign)
            tmp = tree
    return decode_dle_x


# 以 32位 510x512 的 lena 图为例
# 分块
def jpeg_encode(image: cv.Mat):
    h, w, c = image.shape
    decode_image = np.ones((h // 8 * 8, w // 8 * 8, c), dtype=image.dtype)
    # 以高宽8的倍数分块，丢弃小于8x8的子块
    for i in range(h // 8):
        for j in range(w // 8):
            # 以 8x8 的块读取原图像
            block = image[i * 8:(i + 1) * 8, j * 8:(j + 1) * 8]
            # 颜色转换
            block = cv.cvtColor(block, cv.COLOR_BGR2YUV)
            block = block.astype(np.float64)
            # 灰度平移
            block -= 128
            # 分离通道
            y, u, v = cv.split(block)
            # print("原子块Y通道:", y)
            # DCT 变换
            dct_y = cv.dct(y)
            dct_u = cv.dct(u)
            dct_v = cv.dct(v)
            # print("DCT变换结果:", dct_y)
            # 量化
            r_y = np.round(dct_y / L1)
            r_u = np.round(dct_u / L2)
            r_v = np.round(dct_v / L2)
            # print("量化结果:", r_y)
            # 游程编码
            zzag_y = np.zeros(64)
            zzag_u = np.zeros(64)
            zzag_v = np.zeros(64)
            # 先打平
            zzag_y[Z] = r_y
            zzag_u[Z] = r_u
            zzag_v[Z] = r_v
            # 再编码
            dle_y = DLE(zzag_y)
            dle_u = DLE(zzag_u)
            dle_v = DLE(zzag_v)
            # print("游程编码结果:", dle_y)
            # 三通道 huffman 编码
            huffman_y, tree_y = huffman_encode(dle_y)
            huffman_u, tree_u = huffman_encode(dle_u)
            huffman_v, tree_v = huffman_encode(dle_v)
            # print("Huffman编码结果:", huffman_y)
            # 解哈夫曼编码
            decode_dle_y = huffman_decode(huffman_y, tree_y)
            decode_dle_u = huffman_decode(huffman_u, tree_u)
            decode_dle_v = huffman_decode(huffman_v, tree_v)
            # print("Huffman解码结果:", decode_dle_y)
            # 解游程编码
            decode_zzag_y = rDLE(decode_dle_y)
            decode_zzag_u = rDLE(decode_dle_u)
            decode_zzag_v = rDLE(decode_dle_v)
            # 1维转2维
            decode_r_y = decode_zzag_y[Z]
            decode_r_u = decode_zzag_u[Z]
            decode_r_v = decode_zzag_v[Z]
            # print("解游程编码结果:", decode_r_y)

            # 反量化
            decode_dct_y = decode_r_y * L1
            decode_dct_u = decode_r_u * L2
            decode_dct_v = decode_r_v * L2
            # print("反量化结果:", decode_dct_y)
            # 逆 DCT 变换
            decode_y = cv.idct(decode_dct_y)
            decode_u = cv.idct(decode_dct_u)
            decode_v = cv.idct(decode_dct_v)
            # print("逆DCT变换结果:", decode_y)
            # 合并
            decode_block_yuv = cv.merge((decode_y, decode_u, decode_v))
            # 灰度回移
            decode_block_yuv += 128
            decode_block_yuv = decode_block_yuv.astype(np.uint8)
            decode_block_bgr = cv.cvtColor(decode_block_yuv, cv.COLOR_YUV2BGR)
            decode_image[i * 8:(i + 1) * 8, j * 8:(j + 1) * 8] = decode_block_bgr
            # break
        # break
    # decode_image = cv.cvtColor(decode_image, cv.COLOR_YUV2BGR)
    return decode_image


img = cv.imread("lena.png", cv.IMREAD_COLOR)
decode_image = jpeg_encode(img)
cv.imshow('original image', img)
cv.imshow('decode jpeg image', decode_image)
cv.waitKey(0)
cv.destroyAllWindows()
